#include "annot.tree.h"

#include "misc.h"
#include "stream.h"
#include "cache.thread.h"
#include "annot.h"

#include <unordered_map>
#include <QUrl>

struct AnnotLine {
	uint64_t par_id;
	int color_r;
	int color_g;
	int color_b;
	std::string abbrev;
	std::string name;
	int lineno;
};

void AnnotTree::load(std::istream& fs) {
	std::unordered_map<uint64_t, AnnotLine> id2annot;
	std::string line;
	int lineno=0;

	while(fnt_getline(fs, line)) {
		if(line.empty())
			continue;
		if(line[0]=='#')
			continue;

		size_t sp[8];
		char* e;
		uint64_t id=strtoull(&line[0], &e, 10);
		sp[0]=e-&line[0];
		if(sp[0]<=0)
			throwError("Failed to parse ID: ", line);
		if(id<=0)
			throwError("ID should be >0: ", line);
		auto& annot=id2annot[id]={0, 0, 0, 0, {}, {}, lineno++};

		sp[7]=line.size();
		char csep=':';
		if(sp[0]<sp[7]) {
			csep=line[sp[0]];
			sp[0]++;
		}

		for(int i=0; i<3; i++) {
			auto l=line.find(csep, sp[2*i]);
			if(l==std::string::npos) {
				sp[2*i+1]=sp[7];
				sp[2*i+2]=sp[7];
			} else {
				sp[2*i+1]=l;
				sp[2*i+2]=l+1;
			}
		}

		if(sp[0]!=sp[1]) {
			char* e;
			annot.par_id=strtoull(&line[sp[0]], &e, 10);
			if(e!=sp[1]+&line[0])
				throwError("Failed to parse parent ID: ", line);
			if(annot.par_id!=0) {
				if(id2annot.find(annot.par_id)==id2annot.end())
					throwError("Parent ID not previously defined: ", line);
			}
		}
		if(sp[2]!=sp[3]) {
			QColor color{QLatin1String(&line[sp[2]], sp[3]-sp[2])};
			if(!color.isValid())
				throwError("Failed to parse color: ", line);
			annot.color_r=color.red();
			annot.color_g=color.green();
			annot.color_b=color.blue();
		}
		if(sp[4]==sp[5]) {
			if(sp[6]==sp[7]) {
				annot.abbrev=annot.name=std::to_string(id);
			} else {
				annot.abbrev=annot.name=line.substr(sp[6], sp[7]-sp[6]);
			}
		} else {
			if(sp[6]==sp[7]) {
				annot.abbrev=annot.name=line.substr(sp[4], sp[5]-sp[4]);
			} else {
				annot.abbrev=line.substr(sp[4], sp[5]-sp[4]);
				annot.name=line.substr(sp[6], sp[7]-sp[6]);
			}
		}
	}

	if(!fs.eof())
		throwError("Error while reading annot file.");

	if(id2annot.empty())
		throwError("Empty annot file.");

	_items.clear();
	uint32_t i=0;
	uint64_t id=0;
	while(true) {
		typedef decltype(id2annot.begin()) iter;
		std::vector<iter> toremove{};
		int r=0;
		for(auto it=id2annot.begin(); it!=id2annot.end(); ++it) {
			if(it->second.par_id==id) {
				toremove.push_back(it);
			}
		}
		std::sort(toremove.begin(), toremove.end(), [](iter a, iter b) { return a->second.lineno<b->second.lineno; });
		for(auto it: toremove) {
			_items.push_back(AnnotItem{i, it->first, it->second.color_r, it->second.color_g, it->second.color_b, std::move(it->second.abbrev), std::move(it->second.name), 0, 0, r++});
			id2annot.erase(it);
		}

		if(i==0) {
			_root_num=r;
		} else {
			_items[i-1].child_num=r;
		}

		if(i>=_items.size())
			break;

		id=_items[i].id;
		_items[i].child_idx=_items.size()+1;
		i++;
	}

	if(!id2annot.empty())
		throwError("Annot data not conforming to a tree structure.");

	_maps.clear();
	for(uint32_t i=0; i<_items.size(); i++)
		_maps[_items[i].id]=i+1;
}

bool AnnotTree::waitLoading() {
	std::unique_lock<std::mutex> lck{_mtx};
	while(_state==State::Loading)
		_cv.wait(lck);
	return _state==State::Ready?true:false;
}

bool AnnotTree::ensureLoaded() {
	State prevs;
	{
		std::unique_lock<std::mutex> lck{_mtx};
		prevs=_state;
		if(_state==State::Empty)
			_state=State::Loading;
	}

	switch(prevs) {
		case State::Empty:
			try {
				if(isUrl(_loc)) {
					InputStream fs;
					auto cache=CacheThread::instance()->download(_loc);
					if(!cache)
						throwError("Failed to start downloading annot.");
					auto r=CacheThread::instance()->waitForFile(cache);
					if(r!=0 || !cache.isReady())
						throwError("Failed to download annot.");
					if(!fs.open(cache.buffers(), cache.size()))
						throwError("Failed to open annot.");
					load(fs);
				} else {
					InputStream fs;
					if(!fs.open(_loc.toStdString()))
						throwError("Failed to open annot.");
					load(fs);
				}
				std::unique_lock<std::mutex> lck{_mtx};
				_state=State::Ready;
				_cv.notify_all();
			} catch(const std::exception& e) {
				std::unique_lock<std::mutex> lck{_mtx};
				_state=State::Error;
				_cv.notify_all();
				throw e;
			}
			return true;
		case State::Loading:
			return waitLoading();
		case State::Ready:
			return true;
		case State::Error:
			return false;
	}
	return false;
}

template<typename T, typename U> inline void mapDataHelper(const T* dat, U* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride, const std::unordered_map<uint64_t, uint32_t>& idmaps) {
	uint64_t id=0;
	T idmap=0;
	for(int64_t z=0; z<d; z++) {
		for(int64_t y=0; y<h; y++) {
			for(int64_t x=0; x<w; x++) {
				uint64_t v=dat[x+y*w+z*w*h];
				if(v!=id) {
					auto i=idmaps.find(v);
					if(i==idmaps.end()) {
						id=v;
						idmap=0;
					} else {
						id=i->first;
						idmap=i->second;
					}
				}
				cube[(x0+x)+(y0+y)*ystride+(z0+z)*zstride]=idmap;
			}
		}
	}
}
template<typename T> void AnnotTree::mapData(const T* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const {
	auto s=_items.size();
	if(s<=255) {
		mapDataHelper<T, uint8_t>(dat, reinterpret_cast<uint8_t*>(cube), w, h, d, x0, y0, z0, ystride, zstride, _maps);
	} else if(s<=65535) {
		mapDataHelper<T, uint16_t>(dat, reinterpret_cast<uint16_t*>(cube), w, h, d, x0, y0, z0, ystride, zstride, _maps);
	} else {
		mapDataHelper<T, uint32_t>(dat, reinterpret_cast<uint32_t*>(cube), w, h, d, x0, y0, z0, ystride, zstride, _maps);
	}
}
template void AnnotTree::mapData<uint8_t>(const uint8_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<int8_t>(const int8_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<uint16_t>(const uint16_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<int16_t>(const int16_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<uint32_t>(const uint32_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<int32_t>(const int32_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<uint64_t>(const uint64_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;
template void AnnotTree::mapData<int64_t>(const int64_t* dat, char* cube, int64_t w, int64_t h, int64_t d, int64_t x0, int64_t y0, int64_t z0, int64_t ystride, int64_t zstride) const;

AnnotModel::AnnotModel(QObject* s, const AnnotTree& t):
	QAbstractItemModel{s}, tree(t)
{
}

int AnnotModel::columnCount(const QModelIndex& parent) const {
	return 1;
}

QVariant AnnotModel::data(const QModelIndex& index, int role) const {
	if(!index.isValid())
		return QVariant();
	auto idx=index.internalId();
	auto n=tree.annot(idx+1);
	if(!n)
		return QVariant();

	switch(index.column()) {
		case 0:
			{
				switch(role) {
					case Qt::DisplayRole:
						return QString::fromStdString(n->abbrev);
					case Qt::DecorationRole:
						return QColor{n->color_r, n->color_g, n->color_b};
					case Qt::ToolTipRole:
					case Qt::StatusTipRole:
						return QString::fromStdString(n->name);
				}
			}
			break;
	}
	return QVariant();
}

Qt::ItemFlags AnnotModel::flags(const QModelIndex& index) const {
	if(!index.isValid())
		return 0;
	auto idx=index.internalId();
	auto n=tree.annot(idx+1);
	if(!n)
		return 0;
	if(index.column()==0)
		return QAbstractItemModel::flags(index);
	return QAbstractItemModel::flags(index);
}

bool AnnotModel::hasChildren(const QModelIndex& parent) const {
	if(!parent.isValid())
		return tree.root_num()>0;
	auto idx=parent.internalId();
	auto n=tree.annot(idx+1);
	if(!n)
		return false;
	return n->child_num>0;
}

QVariant AnnotModel::headerData(int section, Qt::Orientation orientation, int role) const {
	if(role!=Qt::DisplayRole)
		return QVariant();
	switch(section) {
		case 0:
			return QString("Annotation");
	}
	return QVariant();
}

QModelIndex AnnotModel::index(int row, int column, const QModelIndex& parent) const {
	if(parent.isValid()) {
		auto idx=parent.internalId();
		auto n=tree.annot(idx+1);
		if(!n)
			return QModelIndex();
		if(row>=0 && static_cast<size_t>(row)<n->child_num && column<columnCount()) {
			return createIndex(row, column, n->child_idx+row-1);
		} else {
			return QModelIndex();
		}
	} else {
		if(row>=0 && static_cast<size_t>(row)<tree.root_num() && column<columnCount()) {
			return createIndex(row, column, row);
		} else {
			return QModelIndex();
		}
	}
}

QModelIndex AnnotModel::parent(const QModelIndex& index) const {
	if(!index.isValid())
		return QModelIndex();

	auto idx=index.internalId();
	auto n=tree.annot(idx+1);
	if(!n)
		return QModelIndex();

	auto par_idx=n->par_idx;
	auto pn=tree.annot(par_idx);
	if(!pn)
		return QModelIndex();

	return createIndex(pn->row, 0, par_idx-1);
}

int AnnotModel::rowCount(const QModelIndex& parent) const {
	if(parent.column()>0)
		return 0;
	if(!parent.isValid())
		return tree.root_num();

	auto idx=parent.internalId();
	auto n=tree.annot(idx+1);
	if(!n)
		return 0;
	return n->child_num;
}

bool AnnotModel::setData(const QModelIndex& index, const QVariant& value, int role) {
	return false;
}
QString AnnotModel::getPath(uint32_t id) const {
	std::vector<const AnnotItem*> annots;
	auto n=tree.annot(id);
	while(n) {
		annots.push_back(n);
		n=tree.annot(n->par_idx);
	}
	if(annots.empty())
		return {};
	std::ostringstream oss;
	oss<<annots.back()->abbrev;
	annots.pop_back();
	while(!annots.empty()) {
		oss<<" > ";
		oss<<annots.back()->abbrev;
		annots.pop_back();
	}
	return QString::fromStdString(oss.str());
}

